package dev.langchain4j.example.agent;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;


@Configuration
public class CustomerSupportAgentConfiguration {



    @Bean
    ChatMemoryProvider chatMemoryProvider(TokenCountEstimator tokenizer) {
        return memoryId -> TokenWindowChatMemory.builder()
                .id(memoryId)
                .maxTokens(5000, tokenizer)
                .build();
    }


    // @Bean
    // EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel, ResourceLoader resourceLoader, TokenCountEstimator tokenizer) throws IOException {
    //     // 创建内存向量存储
    //     EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();


    //     //
    //     // 加载示例文档
    //     Resource resource = resourceLoader.getResource("classpath:miles-of-smiles-terms-of-use.txt");
    //     Document document = loadDocument(resource.getFile().toPath(), new TextDocumentParser());
        
    //     // 文档分片和向量化处理
    //     DocumentSplitter documentSplitter = DocumentSplitters.recursive(200, 50, tokenizer);
    //     EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
    //             .documentSplitter(documentSplitter)
    //             .embeddingModel(embeddingModel)
    //             .embeddingStore(embeddingStore)
    //             .build();
    //     ingestor.ingest(document);
        
    //     return embeddingStore;
    // }


    @Bean
    EmbeddingStore<TextSegment> embeddingStore() {
        return MilvusEmbeddingStore.builder()
                .uri("https://in03-d1c03b1f6ba7563.serverless.ali-cn-hangzhou.cloud.zilliz.com.cn")
                .token("6dbc25361baceaa634d8271211479beaedcc63aac6e5bccdeef4cebd7a51dde59628dbfdd535b09ffad400c5ddcfa712d3736f96")
                .collectionName("sports_news")
                .dimension(1024)
                .build();
    }

    @Bean
    ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore, EmbeddingModel embeddingModel) {
        int maxResults = 5;
        double minScore = 0.7;

        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(maxResults)
                .minScore(minScore)
                .build();
    }

    @Bean
    TokenCountEstimator tokenCountEstimator() {
        return new OpenAiTokenCountEstimator(GPT_4_O_MINI);
    }

}
